
import numpy as np
import numpy.fft as fft
from pyfft.cl import Plan
import pyopencl as cl
import pyopencl.array as cl_array
N = 2**3

ctx = cl.create_some_context()
queue = cl.CommandQueue(ctx)

plan = Plan((N, N), queue=queue, dtype=np.complex128, fast_math=False)

a = np.random.rand(N,N)
a_tf = fft.fft2(a)
a_test = fft.ifft2(a_tf)

print 'cpu err', np.abs(a-a_test).max()

a_gpu = cl_array.to_device(queue, a.astype(np.float64))
a_tf_gpu = cl_array.to_device(queue, a.astype(np.complex128))

plan.execute(a_tf_gpu.data)
plan.execute(a_tf_gpu.data, inverse=True)

print 'gpu err', np.abs(a_tf_gpu.get()-a).max()
